package cn.lgwen.spark.ml.learning.kaggle;

import org.apache.spark.ml.classification.RandomForestClassificationModel;
import org.apache.spark.ml.classification.RandomForestClassifier;
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator;
import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.regression.RandomForestRegressionModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

/**
 * 2020/3/16
 * aven.wu
 * danxieai258@163.com
 * 随机森林
 */
public class TitanicRandomForestClass {

    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder().master("local[*]")
                .appName("TitanicRandomForestClass")
                .getOrCreate();

        RandomForestClassifier rf = new RandomForestClassifier()
                .setLabelCol("Survived");

        Dataset<Row> trainData = spark.read().format("csv").option("header", true).option("inferSchema", true)
                .load(TitanicRandomForestClass.class.getResource("/").getPath() + "train.csv");

        VectorAssembler vectorAssem4 = new VectorAssembler()
                .setInputCols(new String[]{"Sex", "SibSp", "Cabin", "Embarked",
                        "Fare", "Age", "Pclass", "Parch"}).
                        setOutputCol("features").setHandleInvalid("keep");
        trainData = vectorAssem4.transform(trainData);
        RandomForestClassificationModel model = rf.fit(trainData);
        System.out.println("Learned classification forest model:\n" + model.toDebugString());

        Dataset<Row> testData = spark.read().format("csv").option("header", true).option("inferSchema", true)
                .load(TitanicRandomForestClass.class.getResource("/").getPath()  + "test.csv");
        testData = vectorAssem4.transform(testData);
        testData = model.transform(testData);

        testData.show();

        BinaryClassificationEvaluator evaluator = new BinaryClassificationEvaluator()
                .setLabelCol("Survived").setRawPredictionCol("prediction");
        double evaluate = evaluator.evaluate(testData);
        System.out.println(evaluate);
    }
}
